package io.github.javpower.imagerex.service;

import lombok.extern.slf4j.Slf4j;
import org.datavec.image.loader.NativeImageLoader;
import org.deeplearning4j.nn.graph.ComputationGraph;
import org.deeplearning4j.nn.transferlearning.TransferLearningHelper;
import org.deeplearning4j.zoo.ZooModel;
import org.deeplearning4j.zoo.model.VGG19;
import org.nd4j.linalg.api.ndarray.INDArray;
import org.nd4j.linalg.dataset.DataSet;
import org.nd4j.linalg.dataset.api.preprocessor.DataNormalization;
import org.nd4j.linalg.dataset.api.preprocessor.VGG16ImagePreProcessor;
import org.springframework.stereotype.Service;

import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.util.ArrayList;
import java.util.List;

@Slf4j
@Service
public class FeatureExtractor {

    private ComputationGraph vgg19;
    private TransferLearningHelper transferLearningHelper;
    private NativeImageLoader loader;
    private DataNormalization scaler;
    public FeatureExtractor() throws IOException {
        // 指定本地模型文件的路径
        File modelFile = new File("models/vgg19_dl4j_inference.zip");
        // 检查模型文件是否存在
        if (modelFile.exists()) {
            vgg19 = ComputationGraph.load(modelFile, true);
        } else {
            ZooModel zooModel = VGG19.builder().build();
            vgg19 = (ComputationGraph) zooModel.initPretrained();
        }
        transferLearningHelper = new TransferLearningHelper(vgg19, "fc2");
        loader = new NativeImageLoader(224, 224, 3);
        scaler = new VGG16ImagePreProcessor();
    }
    public INDArray extractFeatures(String imagePath) throws IOException {
        File file = new File(imagePath);
        try (InputStream inputStream = new FileInputStream(file)) {
            INDArray image = loader.asMatrix(inputStream);
            scaler.transform(image);
            DataSet dataSet = new DataSet(image, null);
            INDArray features = transferLearningHelper.featurize(dataSet).getFeatures();
            return features;
        }
    }

    public List<Float> getVectorFeatures(String path) {
        INDArray indArray = null;
        try {
            indArray = extractFeatures(path);
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
        float[] floatArray = indArray.toFloatVector();
        List<Float> floatList = new ArrayList<>();
        for (float f : floatArray) {
            floatList.add(f);
        }
        return floatList;
    }
}